#include <sys/time.h>
#include <string>
#include <vector>
#include "taiic_avsr.h"
#include "taiic_rknn.h"
#include <dirent.h>
#include <iostream>
#include <fstream>
using namespace std;

/*
rknntoolkit_1.5.0  output_type is int8, need use zp and  scale to trans float
*/
// 获取当前时间
static inline int64_t getCurrentTimeUs()
{
  struct timeval tv;
  gettimeofday(&tv, NULL);
  return tv.tv_sec * 1000000 + tv.tv_usec;
}

string join_paths(const string &dir1, const string &dir2)
{
  std::string full_path = dir1;
  if (full_path.back() != '/')
  {
    full_path += "/";
  }
  full_path += dir2;
  if (full_path.back() != '/')
  {
    full_path += "/";
  }
  return full_path;
}

// 遍历文件夹保存文件名
void EnumerateFilesInDirectory(const std::string &path, std::vector<std::string> &files)
{
  DIR *dir;
  struct dirent *ent;
  if ((dir = opendir(path.c_str())) != NULL)
  {
    while ((ent = readdir(dir)) != NULL)
    {
      // 忽略目录和非文件条目
      if (ent->d_type == DT_REG)
        files.push_back(ent->d_name);
    }
    closedir(dir);
  }
  else
  {
    // 无法打开目录
    perror("");
    return;
  }
}

int main(int argc, char *argv[])
{
  string av_dir = argv[2];
  string video_path = join_paths(av_dir, "video/");
  string audio_path = join_paths(av_dir, "audio/");

  //遍历输入文件夹中的文件并保存文件名
  vector<string> files;
  EnumerateFilesInDirectory(video_path, files);
  cout << "Number of files in the directory: " << files.size() << std::endl;

  // 检查输出文件是否成功打开
  ofstream outfile("result.txt");
  if (!outfile.is_open())
  {
    std::cerr << "Unable to open file result.txt";
    return 1;
  }

  // 模型推理相关初始化
  AVSR_TOOLKIT_MODEL_CTX_S *avsr_ctx;
  avsr_ctx = reinterpret_cast<AVSR_TOOLKIT_MODEL_CTX_S *>(malloc(sizeof(AVSR_TOOLKIT_MODEL_CTX_S))); // 分配内存空间

  // 批量推理
  for (const auto &file : files)
  {
    memset(avsr_ctx, 0, sizeof(AVSR_TOOLKIT_MODEL_CTX_S));
    avsr_ctx->modelPath = argv[1];
    avsr_rknn_toolkit_config_init(avsr_ctx);

    avsr_rknn_toolkit_io_init(avsr_ctx);

    // Get Input Data From CPU
    int i = 0;
    unsigned char *input_data[avsr_ctx->io_num.n_input];
    input_data[i] = new unsigned char[avsr_ctx->input_attrs[i].size]; // 设置输入大小
    FILE *fp = fopen((video_path + file).c_str(), "rb");
    printf("==load video is %s==\n", (video_path + file).c_str());
    if (fp == NULL)
    {
      perror("open failed!");
      return -1;
    }
    fread(input_data[i], avsr_ctx->input_attrs[i].size, 1, fp);
    fclose(fp);
    if (!input_data[i])
      return -1;

    i = 1;
    input_data[i] = new unsigned char[avsr_ctx->input_attrs[i].size]; // 设置输入大小
    FILE *fp1 = fopen((audio_path + file).c_str(), "rb");
    printf("==load audio is %s==\n", (audio_path + file).c_str());
    if (fp1 == NULL)
    {
      perror("open failed!");
      return -1;
    }
    fread(input_data[i], avsr_ctx->input_attrs[i].size, 1, fp1);
    fclose(fp1);
    if (!input_data[i])
      return -1;

    // 模型推理
    avsr_rknn_toolkit_data_refresh(avsr_ctx, input_data);
    int64_t start_us = getCurrentTimeUs();
    int ret = rknn_run(avsr_ctx->context, NULL);
    if (ret < 0)
    {
      RK_LOGE("===rknn_run_fail! ret = %d\n", ret);
      return -1;
    }
    int64_t elapse_us = getCurrentTimeUs() - start_us;
    RK_LOGD(" Elapse Time = %.2fms, FPS = %.2f\n", elapse_us / 1000.f, 1000.f * 1000.f / elapse_us);
    //MODEL_RESULT_S rknn_result = avsr_rknn_toolkit_result_int8(avsr_ctx);
    MODEL_RESULT_S rknn_result = avsr_rknn_toolkit_result_int8_opt(avsr_ctx);
    printf("===result label is %d, prob is %f====\n", rknn_result.label, rknn_result.prob);

    // Destroy rknn memory
    for (uint32_t i = 0; i < avsr_ctx->io_num.n_input; ++i)
    {
      rknn_destroy_mem(avsr_ctx->context, avsr_ctx->input_mems[i]);
    }
    for (uint32_t i = 0; i < avsr_ctx->io_num.n_output; ++i)
    {
      rknn_destroy_mem(avsr_ctx->context, avsr_ctx->output_mems[i]);
    }
    // destroy
    rknn_destroy(avsr_ctx->context);
    for (int i = 0; i < avsr_ctx->io_num.n_input; i++)
      free(input_data[i]);

    // 结果写入文件
    outfile << file << ": result label: " << rknn_result.label << ", prob: " << rknn_result.prob << std::endl;
  }

  // 关闭文件
  outfile.close();

  return 0;
}
